/*
题目: 移除子树后的二叉树高度
给你一棵 二叉树 的根节点 root ，树中有 n 个节点。每个节点都可以被分配一个从 1 到 n 且互不相同的值。另给你一个长度为 m 的数组 queries 。

你必须在树上执行 m 个 独立 的查询，其中第 i 个查询你需要执行以下操作：

    从树中 移除 以 queries[i] 的值作为根节点的子树。题目所用测试用例保证 queries[i] 不 等于根节点的值。

返回一个长度为 m 的数组 answer ，其中 answer[i] 是执行第 i 个查询后树的高度。

注意：

    查询之间是独立的，所以在每个查询执行后，树会回到其 初始 状态。
    树的高度是从根到树中某个节点的 最长简单路径中的边数 。

https://leetcode.cn/problems/height-of-binary-tree-after-subtree-removal-queries
 */
public class TreeQueries {
    int[] src = null;
    // 以 root 为树根时 的高度
    Map<TreeNode, Integer> treeHeight = new HashMap<> ();

    public int[] treeQueries(TreeNode root, int[] queries) {
        getHeight(root); // 获取所有子树的高度

        treeHeight.put(null, 0); // 不用写 getOrDefault
        src = new int[treeHeight.size()];

        // 获取 去掉 root 之后, 剩下的树的高度
        // 由于题目, 所以高度要从 -1 开始算, 去掉源根节点后, restHeight = 0
        dfs(root, -1, 0);

        // 然后再将提取出 要去掉的结点之后 的树的高度即可
        for (int i = 0; i < queries.length; i ++) {
            queries[i] = src[queries[i]]; // 去掉这个结点后的高度
        }
        return queries;
    }

    private void dfs(TreeNode root, int curHeight, int restHeight) {
        if (root == null) return ;
        ++ curHeight; // 高度增加

        // 去掉当前结点后  高度就是 restHeight
        src[root.val] = restHeight;
        // restHeight : 去掉当前结点后, 「整棵树」的高度
        // 每一次 dfs, 我们尝试去掉 root, 那么能比较的就是 「当前另一边的树的高度」 和 「整棵树另一边的树的高度」
        dfs(root.left, curHeight, Math.max(restHeight, curHeight + treeHeight.get(root.right)));
        dfs(root.right, curHeight, Math.max(restHeight, curHeight + treeHeight.get(root.left)));
    }

    // 获取以 root 为树根的树的 高度
    private int getHeight(TreeNode root) {
        if (root == null)   return 0;
        int h = Math.max(getHeight(root.left), getHeight(root.right)) + 1;
        treeHeight.put(root, h);
        return h;
    }
}
